'''
Created on Jan 16, 2010

@author: oabalbin
'''

import numpy as np
import scipy as sp
from scipy import stats
import rpy2.robjects as robjects
import rpy2.robjects.numpy2ri
from rpy2.robjects.packages import importr

import signatures.parsers.read_microarray_data as pm

class bfrm_parser:
    
    def __init__(self, inter, covar, nresfac, nlatfac, ppith, nvar, nfactors):
        """
        inter=bool, covar=covariates, 
        nresfac = number of response factors, nlatfac = number of latent factors.
        ppith threshold for the posterior probability of inclusion.
        nvar=number of variables.
        nfactors=number of final factors in the file mA.txt
        """
        self.inter = inter
        self.covar = covar
        self.nresfac = nresfac
        self.nlatfac = nlatfac
        self.ppith = ppith
        self.nvar = nvar
        self.nfactors = nfactors
 
        
    def get_var_from_mPostPib(self, mPostPib):
        """
        It gets the variables with inclusions prob > ppith from the mPostPib.txt bfrm output file. 
        This file contains the posterior means for inclusion probabilities. 
        rows= X, Y variables in dataset, columns: inter=intercept, covar=covariates, 
        nresfac = number of response factors, nlatfac = number of latent factors.
        ppith threshold for the posterior probability of inclusion. 
        """
        
        # loading the mPostPib matrix
        #m = np.fromfile(file, dtype=float, sep='\t')
        #mPostPib = np.reshape(m, (self.nvar,-1))
        
        # True False indicators which can be used to select the elements of another matrix, for example mA
        varInd = mPostPib > self.ppith
        
        return varInd 
        
    def get_var_loadings(self, mA, varInd):
        """
        Input: 
        The mA.txt file generated in the bfrm output
        varInd = Boolean Indicators for each variable in each factor. True means that a variable has inclusion prob > ppith in that factor.
        varInd is the output of the function get_var_from_mPostPib
        
        It returns two list of numpy arrays.
        varbyfact = It has the index for each variable with inclusion prob > ppith in each of the factors. Each factor is represented by a numpy array
        loadbyfact = It has the loading for each variable with inclusion prob > ppith in each of the factors. Each factor is represented by a numpy array
        varbyfact and loadbyfact correspond on to one.
        """
        #  the mA matrix with the loadings for each variable.
        #lf = np.fromfile(file, dtype=float, sep='\t')
        #mA = np.reshape(lf, (self.nvar,-1))
        
        fact = range(0,mA.shape[1]) 
        loadbyfact= []
        varbyfact = []
        # p[h[:,1],1], # ft[h[:,2]]
        vars = np.array(range(0,self.nvar)) 
        
        for ft in fact:
            loading = mA[varInd[:,ft],ft]
            loadbyfact.append(loading)
            varth = vars[varInd[:,ft]]
            varbyfact.append(varth)
            #print varth
            #print loading
        
        #print varbyfact #, loadbyfact
        return varbyfact, loadbyfact
    
        
    def get_gene_in_module(self, outfile, mPostPib, mA, genelist, geneseed, mVarIn):
        """
        This function calculates a "pathway" activity as total weighted expression of the genes in a module.
        The genes are weigthed based on their contribution (load) to each of the factors
        patwactivity = (exp of gen i in factor j)*(load of gen in factor j) for each cell line k
        input: output file, mPostPibFile, mAfile, genelist in the experiment, geneseed list in the experiment
        mVarIn: order in which the variables were selected in the final model by BFRM
        """
        # Loading the data
        '''
        lf = np.fromfile(mAfile, dtype=float, sep='\t')
        mA = np.reshape(lf, (self.nvar,-1))
        
        m = np.fromfile(mPostPibFile, dtype=float, sep='\t')
        mPostPib = np.reshape(m, (self.nvar,-1))
        '''
        
        #print mPostPibFile, mAfile
        varInd = self.get_var_from_mPostPib(mPostPib)
        varbyfact, loadbyfact = self.get_var_loadings(mA, varInd)
        geneinfactor={}
        
        
        nf=0
        for v,l in zip(varbyfact, loadbyfact):
            print '##############'+'\t'+'Factor number' + '\t'+ str(nf)
            outfile.write('##############'+'\t'+'Factor number' + '\t'+ str(nf)+'\n')
            for g,q in zip(list(v),list(l)):
                p = mVarIn[g] 
                gen = genelist[p]
                if gen in geneseed:
                    print str(p)+'\t'+ gen +'\t'+ str(q) + '\t' + '#### seed gene ####'
                    outfile.write(str(p)+'\t'+ gen +'\t'+ str(q) + '\t' + '#### seed gene ####'+'\n')
                else:
                    print str(p)+'\t'+ gen +'\t'+ str(q)
                    outfile.write(str(p)+'\t'+ gen +'\t'+ str(q)+'\n')
                
                geneinfactor[gen] = p
            
            nf+=1
        
        outfile.write('##############'+'\t'+'Factor number' + '\t'+ str(00)+'\n')
        
        return geneinfactor
    
    def get_gene_in_module_uniq(self, outfile, mPostPib, mA, genelist, geneseed, mVarIn, Psi, varVariance):
        """
        This function calculates a "pathway" activity as total weighted expression of the genes in a module.
        The genes are weigthed based on their contribution (load) to each of the factors
        patwactivity = (exp of gen i in factor j)*(load of gen in factor j) for each cell line k
        input: output file, mPostPibFile, mAfile, genelist in the experiment, geneseed list in the experiment
        mVarIn: order in which the variables were selected in the final model by BFRM
        """

        #print mPostPibFile, mAfile
        varInd = self.get_var_from_mPostPib(mPostPib)
        varbyfact, loadbyfact = self.get_var_loadings(mA, varInd)
        geneinfactor={}
        

        nf=0
        header =['genInd','geneName','load','commonality','uniqueness','variance','commonality_raw','uniqueness_raw','seed']
        outfile.write(",".join(header).replace(',','\t')+'\n')
        for v,l in zip(varbyfact, loadbyfact):
            print '##############'+'\t'+'Factor number' + '\t'+ str(nf)
            outfile.write('##############'+'\t'+'Factor number' + '\t'+ str(nf)+'\n')
            for g,q in zip(list(v),list(l)):
                p = mVarIn[g] 
                gen = genelist[p]
                uniqueness2 = Psi[g]
                variance = varVariance[g,0]
                uniqueness = varVariance[g,2]/varVariance[g,0]
                commonality = 1.0 - varVariance[g,2]/varVariance[g,0]
                commonality2 = varVariance[g,1]
                
                if gen in geneseed:
                    record = [str(p),gen,str(q),str(commonality), str(uniqueness), str(variance),str(commonality2),str(uniqueness2),'#### seed gene ####']
                    print ",".join(record).replace(',','\t')
                    outfile.write(",".join(record).replace(',','\t')+'\n')
                else:
                    record = [str(p),gen,str(q),str(commonality), str(uniqueness), str(variance),str(commonality2),str(uniqueness2)]
                    print ",".join(record).replace(',','\t')
                    outfile.write(",".join(record).replace(',','\t')+'\n')
                
                geneinfactor[gen] = p
            
            nf+=1
        
        outfile.write('##############'+'\t'+'Factor number' + '\t'+ str(00)+'\n')
        
        return geneinfactor

    
    
    def get_gene_in_thismodule(self, mPostPibFile, mAfile, genelist, mastergenelist, mVarIn, selection=False):
        """
        This function calculates a "pathway" activity as total weighted expression of the genes in a module.
        The genes are weigthed based on their contribution (load) to each of the factors
        patwactivity = (exp of gen i in factor j)*(load of gen in factor j) for each cell line k
        input: output file, mPostPibFile, mAfile, genelist in the experiment, mastergenelist a list with all genes in the experiments
        """
        # threshold for the highest loadings in the factor
        percent =90
        #  the mA matrix with the loadings for each variable.
        lf = np.fromfile(mAfile, dtype=float, sep='\t')
        mA = np.reshape(lf, (self.nvar,-1))
                
        #print mPostPibFile, mAfile
        varInd = self.get_var_from_mPostPib(mPostPibFile)
        varbyfact, loadbyfact = self.get_var_loadings(mAfile, varInd)
        
        new_mA = np.zeros((len(mastergenelist),mA.shape[1]))
        
        nf=0
        for v,l in zip(varbyfact, loadbyfact):
            if selection:
                loadth = sp.stats.scoreatpercentile(l,percent)
                
            for g,q in zip(list(v),list(l)):
                p = mVarIn[g]
                gen = genelist[p]
                
                if gen in mastergenelist:
                    if selection:
                        if np.abs(q) >= loadth:
                            i = mastergenelist.index(gen)
                            new_mA[i,nf]= q
                        else:
                            i = mastergenelist.index(gen)
                            new_mA[i,nf]= 0.0                        
                    else:
                        i = mastergenelist.index(gen)
                        new_mA[i,nf]= q
                        
            nf+=1
        
        
        return new_mA


    def calc_factorvar(self, mA):
        """
        Returns the variance explained for each factor
        """
        factorvar = (np.sum(np.power(mA,2),axis=0)/mA.shape[0])*100
        return factorvar
    
    def calc_varcommunality(self,expVals,mPsi,varInd,nfactors):
        """
        Returns the communality of each variable. Not completly working because the value of the loadings are not between 0-1 yet.
        """
        # calculate variance
        gene_variance = np.var(expVals[varInd,:],axis=1)
        # var(x) = commonality - uniqueness
        commonality = gene_variance - mPsi
        commonality = np.reshape(commonality,(commonality.shape[0],-1))
        gene_variance = np.reshape(gene_variance,(gene_variance.shape[0],-1))
        mPsi = np.reshape(mPsi,(mPsi.shape[0],-1))
        varcommon = np.concatenate((gene_variance,commonality,mPsi),axis=1)
        
        return varcommon
             
    
    def calc_factorscore(self,mA, expMatval):
        
        facscore = np.dot(mA.T,expMatval)
        return facscore 
    
    def calc_factorscore_bfrm(self,mA, mPostPib, expMatval, mPsi, hvec, Bmat=np.array([]), vec=False):
        """
        It projects a the factors obtained using bfrm into a new expMat matrix.
        It does it following the formula present it in Lucas et al PLoS ONE 2009. voll4, e4523
        Lnew = (Id + A'*Y^-1 *A)*A'Y^-1*(Xnew - B*hnew)
        Lnew = new factor scores, I: Identity, A matrix of loadings, Y: matrix of variance error, Xnew: new matrix of expression values: 
        B: matrix of regression values, not present if running BFRM for latent factors. In this case Bmat correspond to mean of the 
        expression values for each gene across al samples. Bmat (pxn)
        h: vector (or matrix) of design-correction factors, H(nxd). The function assumess that H' was passed as argument
         If not correction due to particular correlation, h is usually a vector of 1
        Note that in this cases mPostPib was not used, but it could be used to preselect the genes with significant loadings.  
        """
                    
        mPsi = np.identity(mPsi.shape[0])*np.transpose(mPsi)
        
        #Lnew = (Id + A'*Y^-1 *A)*A'Y^-1*(Xnew - B*hnew)
        # Weight the mA matrix according to the probability of being 0
        mA = np.multiply(mA,mPostPib)
        # Calculate A'Y^-1
        mATYmat = np.dot(mA.T,np.linalg.inv(mPsi))
        # Calculate A'Y^-1 A
        mATYmatmA = np.dot(mATYmat,mA)
        # Calculate A'Y + A'Y^-1 A
        rhs1 = np.add(mATYmat,np.dot(mATYmatmA,mATYmat)) 
        
        # Calculate Bhnew
        if not Bmat.any():
            Bmat=np.ones((hvec.shape[0],1))
            # It is possible that you need to reshape hvec np.reshape((expMatval.shape[0],-1))
            b= Bmat*hvec
            Bhnew = np.tile(b,expMatval.shape[1])
        elif vec:
            Bmat = np.reshape(Bmat,(Bmat.shape[0],-1))
            Bhnew = Bmat*hvec
        else:
            Bhnew = Bmat*hvec.T
            #print Bhnew
            #print Bhnew.shape
        
        #Calculate Znew
        rhs2 =  expMatval - Bhnew
        # Calculate Lnew = (A'Y^-1 + A'*Y^-1 *A*A'Y^-1)*Znew
        Lnew = np.dot(rhs1,rhs2)
        #Lnew = np.dot(mATYmat,expMatval)
        
        return Lnew 
    
    def rotate_mA(self, mA, rottype, mPostPib=[]):
        """
        Performs rotation of matrix loadings.
        Input: mA matrix of factor loadings.
        Output: rotmA rotate matrix of factor loadings. 
        rottype: type of roation performed, varimax, orthogonal, oblique, or oblique partially specified target rotation 
        pstQ is not working yet: Problem ==> how to specify target
        """
        
        rottype
        r = robjects.r
        r('''library(GPArotation)''')
        
        if rottype=='varimax':
            rotmA = r.varimax(mA)
        elif rottype=='promax':
            rotmA = r.promax(mA)
        elif rottype=='oblique':
            rotmA = r.oblimin(mA)
        elif rottype=='pstQ' and mPostPib.shape[0] > 0:
            rotmA = r.pstQ(mA,W=mPostPib)
        
        rotmAload = np.array(r.loadings(rotmA))
        #print rotmA
        
        return rotmA, rotmAload






'''
mA = np.loadtxt("/home/oabalbin/projects/networks/output/2010_03_08_14_28/bfrm/mA.txt",dtype=float, delimiter='\t')
#mA = np.loadtxt("/home/oabalbin/projects/networks/bfrm_examples/ex2/mA.txt",dtype=float, delimiter='\t')
mPostPib = np.loadtxt("/home/oabalbin/projects/networks/output/2010_03_08_14_28/bfrm/mPostPib.txt",dtype=float, delimiter='\t')
mPsi=np.loadtxt("/home/oabalbin/projects/networks/output/2010_03_08_14_28/bfrm/mPsi.txt",dtype=float, delimiter='\t')
expVal = np.loadtxt("/home/oabalbin/projects/networks/output/2010_03_08_14_28/sdv/2010_03_08_14_28_mat.sdv.test2.txt",dtype=float, delimiter='\t')

outfile5 = "/home/oabalbin/projects/networks/output/2010_03_08_14_28/bfrm/factortest.txt"
bf = bfrm_parser(False, 0 ,0, 0, 0.99,mA.shape[0],mA.shape[1])
#bf.calc_factorvar(mA[:,1:])
#bf.calc_varcommunality(mA)
#sys.exit(0)


# calc_factorscore_bfrm(self,mA, mPostPib, expMatval, mPsi, hvec, Bmat=[])
outfile=open(outfile5,'w')
print mPsi
print mPsi.shape[0]
print np.transpose(mPsi).shape
factscore = bf.calc_factorscore_bfrm( mA[:,1:], mPostPib[:,1:],expVal,mPsi,np.ones((mPsi.shape[0],1)) )
print factscore


for i in range(factscore.shape[0]):
    le = list(factscore[i,:])
    outfile.write(",".join(map(str,le)).replace(',','\t')+'\n')
 
rotmA, rotmAload = bf.rotate_mA(mA[:,1:],'oblique')
'''

'''
mA = np.loadtxt("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/mA.txt",dtype=float, delimiter='\t')
mVarin = np.loadtxt("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/mVariablesIn.txt",dtype=int, delimiter='\t')
mPostPib = np.loadtxt("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/mPostPib.txt",dtype=float, delimiter='\t')
mPsi=np.loadtxt("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/mPsi.txt",dtype=float, delimiter='\t')
expVals = np.loadtxt("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/2010_04_21_10_31_dataset.mat.txt",dtype=float, delimiter='\t')
outfile5 = "/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/2010_04_21_10_31_geneInModules_commonality_all.txt"
tp = pm.parser()
genelist = tp.list_of_names(open("/home/oabalbin/projects/networks/output/2010_04_21_10_31/bfrm/2010_04_21_10_31_dataset.gen.txt"))
geneseed=['ERG','PRKDC','PARP1','DDX5','AR','NUMA1','PRPF8','TOP2B','XRCC5','ABL1','ACP5','ACTL6A']
#print rotmAload
bf = bfrm_parser(False, 0 ,0, 0, 0.99,mA.shape[0],mA.shape[1])
mVarin=mVarin-1
varcommon = bf.calc_varcommunality(expVals,mPsi,mVarin,mA.shape[1])
geneinfactor = bf.get_gene_in_module_uniq(open(outfile5,'w'), mPostPib[:,0:], mA[:,0:], genelist, geneseed,mVarin,mPsi,varcommon)
'''
